#pragma once
#include "Graph.h"
#include "imnodes.h"
#include "imgui.h"
#include "KeyBoard.h"
#include <algorithm>
#include <cassert>
#include <chrono>
#include <cmath>
#include <vector>

namespace Editor
{
	
		enum class NodeType
		{
			add,
			multiply,
			output,
			sine,
			time,
			value
		};

		struct Node
		{
			NodeType type;
			float    value;

			explicit Node(const NodeType t) : type(t), value(0.f) {}

			Node(const NodeType t, const float v) : type(t), value(v) {}
		};

		static bool  emulate_three_button_mouse = false;
		
		class ColorNodeEditor
		{
		public:
			ColorNodeEditor() : graph_(), nodes_(), root_node_id_(-1),
				minimap_location_(ImNodesMiniMapLocation_BottomRight) 
			{
				current_time_seconds = 0.0;
			}
			void Init() 
			{
				ImNodesIO& io = ImNodes::GetIO();
				io.LinkDetachWithModifierClick.Modifier = &ImGui::GetIO().KeyCtrl;
			}
			void show(float totalTime)
			{
				// Update timer context
				current_time_seconds = totalTime;

				auto flags = ImGuiWindowFlags_MenuBar;

				// The node editor window
				ImGui::Begin("color node editor", NULL, flags);

				if (ImGui::BeginMenuBar())
				{
					if (ImGui::BeginMenu("Mini-map"))
					{
						const char* names[] = {
							"Top Left",
							"Top Right",
							"Bottom Left",
							"Bottom Right",
						};
						int locations[] = {
							ImNodesMiniMapLocation_TopLeft,
							ImNodesMiniMapLocation_TopRight,
							ImNodesMiniMapLocation_BottomLeft,
							ImNodesMiniMapLocation_BottomRight,
						};

						for (int i = 0; i < 4; i++)
						{
							bool selected = minimap_location_ == locations[i];
							if (ImGui::MenuItem(names[i], NULL, &selected))
								minimap_location_ = locations[i];
						}
						ImGui::EndMenu();
					}

					ImGui::EndMenuBar();
				}

				ImGui::Columns(2);
				ImGui::TextUnformatted("Right Button -- add node");
				ImGui::TextUnformatted("Delete -- delete selected node or link");
				ImGui::NextColumn();
				if (ImGui::Checkbox("emulate_three_button_mouse", &emulate_three_button_mouse))
				{
					ImNodes::GetIO().EmulateThreeButtonMouse.Modifier =
						emulate_three_button_mouse ? &ImGui::GetIO().KeyAlt : NULL;
				}
				ImGui::Columns(1);

				ImNodes::BeginNodeEditor();

				// Handle new nodes
				// These are driven by the user, so we place this code before rendering the nodes
				{
					const bool open_popup = ImGui::IsWindowFocused(ImGuiFocusedFlags_RootAndChildWindows) &&
						ImNodes::IsEditorHovered() &&
						ImGui::IsMouseDown(1);
						

					ImGui::PushStyleVar(ImGuiStyleVar_WindowPadding, ImVec2(8.f, 8.f));
					if (!ImGui::IsAnyItemHovered() && open_popup)
					{
						ImGui::OpenPopup("add node");
					}

					if (ImGui::BeginPopup("add node"))
					{
						const ImVec2 click_pos = ImGui::GetMousePosOnOpeningCurrentPopup();

						if (ImGui::MenuItem("add"))
						{
							const Node value(NodeType::value, 0.f);
							const Node op(NodeType::add);

							UiNode ui_node;
							ui_node.type = UiNodeType::add;
							ui_node.add.lhs = graph_.insert_node(value);
							ui_node.add.rhs = graph_.insert_node(value);
							ui_node.id = graph_.insert_node(op);

							graph_.insert_edge(ui_node.id, ui_node.add.lhs);
							graph_.insert_edge(ui_node.id, ui_node.add.rhs);

							nodes_.push_back(ui_node);
							ImNodes::SetNodeScreenSpacePos(ui_node.id, click_pos);
						}

						if (ImGui::MenuItem("multiply"))
						{
							const Node value(NodeType::value, 0.f);
							const Node op(NodeType::multiply);

							UiNode ui_node;
							ui_node.type = UiNodeType::multiply;
							ui_node.multiply.lhs = graph_.insert_node(value);
							ui_node.multiply.rhs = graph_.insert_node(value);
							ui_node.id = graph_.insert_node(op);

							graph_.insert_edge(ui_node.id, ui_node.multiply.lhs);
							graph_.insert_edge(ui_node.id, ui_node.multiply.rhs);

							nodes_.push_back(ui_node);
							ImNodes::SetNodeScreenSpacePos(ui_node.id, click_pos);
						}

						if (ImGui::MenuItem("output") && root_node_id_ == -1)
						{
							const Node value(NodeType::value, 0.f);
							const Node out(NodeType::output);

							UiNode ui_node;
							ui_node.type = UiNodeType::output;
							ui_node.output.r = graph_.insert_node(value);
							ui_node.output.g = graph_.insert_node(value);
							ui_node.output.b = graph_.insert_node(value);
							ui_node.id = graph_.insert_node(out);

							graph_.insert_edge(ui_node.id, ui_node.output.r);
							graph_.insert_edge(ui_node.id, ui_node.output.g);
							graph_.insert_edge(ui_node.id, ui_node.output.b);

							nodes_.push_back(ui_node);
							ImNodes::SetNodeScreenSpacePos(ui_node.id, click_pos);
							root_node_id_ = ui_node.id;
						}

						if (ImGui::MenuItem("sine"))
						{
							const Node value(NodeType::value, 0.f);
							const Node op(NodeType::sine);

							UiNode ui_node;
							ui_node.type = UiNodeType::sine;
							ui_node.sine.input = graph_.insert_node(value);
							ui_node.id = graph_.insert_node(op);

							graph_.insert_edge(ui_node.id, ui_node.sine.input);

							nodes_.push_back(ui_node);
							ImNodes::SetNodeScreenSpacePos(ui_node.id, click_pos);
						}

						if (ImGui::MenuItem("time"))
						{
							UiNode ui_node;
							ui_node.type = UiNodeType::time;
							ui_node.id = graph_.insert_node(Node(NodeType::time));

							nodes_.push_back(ui_node);
							ImNodes::SetNodeScreenSpacePos(ui_node.id, click_pos);
						}

						ImGui::EndPopup();
					}
					ImGui::PopStyleVar();
				}

				for (const UiNode& node : nodes_)
				{
					switch (node.type)
					{
					case UiNodeType::add:
					{
						const float node_width = 100.f;
						ImNodes::BeginNode(node.id);

						ImNodes::BeginNodeTitleBar();
						ImGui::TextUnformatted("add");
						ImNodes::EndNodeTitleBar();
						{
							ImNodes::BeginInputAttribute(node.add.lhs);
							const float label_width = ImGui::CalcTextSize("left").x;
							ImGui::TextUnformatted("left");
							if (graph_.num_edges_from_node(node.add.lhs) == 0ull)
							{
								ImGui::SameLine();
								ImGui::PushItemWidth(node_width - label_width);
								ImGui::DragFloat("##hidelabel", &graph_.node(node.add.lhs).value, 0.01f);
								ImGui::PopItemWidth();
							}
							ImNodes::EndInputAttribute();
						}

						{
							ImNodes::BeginInputAttribute(node.add.rhs);
							const float label_width = ImGui::CalcTextSize("right").x;
							ImGui::TextUnformatted("right");
							if (graph_.num_edges_from_node(node.add.rhs) == 0ull)
							{
								ImGui::SameLine();
								ImGui::PushItemWidth(node_width - label_width);
								ImGui::DragFloat("##hidelabel", &graph_.node(node.add.rhs).value, 0.01f);
								ImGui::PopItemWidth();
							}
							ImNodes::EndInputAttribute();
						}

						ImGui::Spacing();

						{
							ImNodes::BeginOutputAttribute(node.id);
							const float label_width = ImGui::CalcTextSize("result").x;
							ImGui::Indent(node_width - label_width);
							ImGui::TextUnformatted("result");
							ImNodes::EndOutputAttribute();
						}

						ImNodes::EndNode();
					}
					break;
					case UiNodeType::multiply:
					{
						const float node_width = 100.0f;
						ImNodes::BeginNode(node.id);

						ImNodes::BeginNodeTitleBar();
						ImGui::TextUnformatted("multiply");
						ImNodes::EndNodeTitleBar();

						{
							ImNodes::BeginInputAttribute(node.multiply.lhs);
							const float label_width = ImGui::CalcTextSize("left").x;
							ImGui::TextUnformatted("left");
							if (graph_.num_edges_from_node(node.multiply.lhs) == 0ull)
							{
								ImGui::SameLine();
								ImGui::PushItemWidth(node_width - label_width);
								ImGui::DragFloat(
									"##hidelabel", &graph_.node(node.multiply.lhs).value, 0.01f);
								ImGui::PopItemWidth();
							}
							ImNodes::EndInputAttribute();
						}

						{
							ImNodes::BeginInputAttribute(node.multiply.rhs);
							const float label_width = ImGui::CalcTextSize("right").x;
							ImGui::TextUnformatted("right");
							if (graph_.num_edges_from_node(node.multiply.rhs) == 0ull)
							{
								ImGui::SameLine();
								ImGui::PushItemWidth(node_width - label_width);
								ImGui::DragFloat(
									"##hidelabel", &graph_.node(node.multiply.rhs).value, 0.01f);
								ImGui::PopItemWidth();
							}
							ImNodes::EndInputAttribute();
						}

						ImGui::Spacing();

						{
							ImNodes::BeginOutputAttribute(node.id);
							const float label_width = ImGui::CalcTextSize("result").x;
							ImGui::Indent(node_width - label_width);
							ImGui::TextUnformatted("result");
							ImNodes::EndOutputAttribute();
						}

						ImNodes::EndNode();
					}
					break;
					case UiNodeType::output:
					{
						const float node_width = 100.0f;
						ImNodes::PushColorStyle(ImNodesCol_TitleBar, IM_COL32(11, 109, 191, 255));
						ImNodes::PushColorStyle(ImNodesCol_TitleBarHovered, IM_COL32(45, 126, 194, 255));
						ImNodes::PushColorStyle(ImNodesCol_TitleBarSelected, IM_COL32(81, 148, 204, 255));
						ImNodes::BeginNode(node.id);

						ImNodes::BeginNodeTitleBar();
						ImGui::TextUnformatted("output");
						ImNodes::EndNodeTitleBar();

						ImGui::Dummy(ImVec2(node_width, 0.f));
						{
							ImNodes::BeginInputAttribute(node.output.r);
							const float label_width = ImGui::CalcTextSize("r").x;
							ImGui::TextUnformatted("r");
							if (graph_.num_edges_from_node(node.output.r) == 0ull)
							{
								ImGui::SameLine();
								ImGui::PushItemWidth(node_width - label_width);
								ImGui::DragFloat(
									"##hidelabel", &graph_.node(node.output.r).value, 0.01f, 0.f, 1.0f);
								ImGui::PopItemWidth();
							}
							ImNodes::EndInputAttribute();
						}

						ImGui::Spacing();

						{
							ImNodes::BeginInputAttribute(node.output.g);
							const float label_width = ImGui::CalcTextSize("g").x;
							ImGui::TextUnformatted("g");
							if (graph_.num_edges_from_node(node.output.g) == 0ull)
							{
								ImGui::SameLine();
								ImGui::PushItemWidth(node_width - label_width);
								ImGui::DragFloat(
									"##hidelabel", &graph_.node(node.output.g).value, 0.01f, 0.f, 1.f);
								ImGui::PopItemWidth();
							}
							ImNodes::EndInputAttribute();
						}

						ImGui::Spacing();

						{
							ImNodes::BeginInputAttribute(node.output.b);
							const float label_width = ImGui::CalcTextSize("b").x;
							ImGui::TextUnformatted("b");
							if (graph_.num_edges_from_node(node.output.b) == 0ull)
							{
								ImGui::SameLine();
								ImGui::PushItemWidth(node_width - label_width);
								ImGui::DragFloat(
									"##hidelabel", &graph_.node(node.output.b).value, 0.01f, 0.f, 1.0f);
								ImGui::PopItemWidth();
							}
							ImNodes::EndInputAttribute();
						}
						ImNodes::EndNode();
						ImNodes::PopColorStyle();
						ImNodes::PopColorStyle();
						ImNodes::PopColorStyle();
					}
					break;
					case UiNodeType::sine:
					{
						const float node_width = 100.0f;
						ImNodes::BeginNode(node.id);

						ImNodes::BeginNodeTitleBar();
						ImGui::TextUnformatted("sine");
						ImNodes::EndNodeTitleBar();

						{
							ImNodes::BeginInputAttribute(node.sine.input);
							const float label_width = ImGui::CalcTextSize("number").x;
							ImGui::TextUnformatted("number");
							if (graph_.num_edges_from_node(node.sine.input) == 0ull)
							{
								ImGui::SameLine();
								ImGui::PushItemWidth(node_width - label_width);
								ImGui::DragFloat(
									"##hidelabel", &graph_.node(node.sine.input).value, 0.01f, 0.f, 1.0f);
								ImGui::PopItemWidth();
							}
							ImNodes::EndInputAttribute();
						}

						ImGui::Spacing();

						{
							ImNodes::BeginOutputAttribute(node.id);
							const float label_width = ImGui::CalcTextSize("output").x;
							ImGui::Indent(node_width - label_width);
							ImGui::TextUnformatted("output");
							ImNodes::EndInputAttribute();
						}

						ImNodes::EndNode();
					}
					break;
					case UiNodeType::time:
					{
						ImNodes::BeginNode(node.id);

						ImNodes::BeginNodeTitleBar();
						ImGui::TextUnformatted("time");
						ImNodes::EndNodeTitleBar();

						ImNodes::BeginOutputAttribute(node.id);
						ImGui::Text("output");
						ImNodes::EndOutputAttribute();

						ImNodes::EndNode();
					}
					break;
					}
				}

				for (const auto& edge : graph_.edges())
				{
					// If edge doesn't start at value, then it's an internal edge, i.e.
					// an edge which links a node's operation to its input. We don't
					// want to render node internals with visible links.
					if (graph_.node(edge.from).type != NodeType::value)
						continue;

					ImNodes::Link(edge.id, edge.from, edge.to);
				}

				ImNodes::MiniMap(0.2f, minimap_location_);
				ImNodes::EndNodeEditor();

				// Handle new links
				// These are driven by Imnodes, so we place the code after EndNodeEditor().

				{
					int start_attr, end_attr;
					if (ImNodes::IsLinkCreated(&start_attr, &end_attr))
					{
						const NodeType start_type = graph_.node(start_attr).type;
						const NodeType end_type = graph_.node(end_attr).type;

						const bool valid_link = start_type != end_type;
						if (valid_link)
						{
							// Ensure the edge is always directed from the value to
							// whatever produces the value
							if (start_type != NodeType::value)
							{
								std::swap(start_attr, end_attr);
							}
							graph_.insert_edge(start_attr, end_attr);
						}
					}
				}

				// Handle deleted links

				{
					int link_id;
					if (ImNodes::IsLinkDestroyed(&link_id))
					{
						graph_.erase_edge(link_id);
					}
				}

				{
					const int num_selected = ImNodes::NumSelectedLinks();
					if (num_selected > 0 && ImGui::IsKeyReleased(DirectX::Keyboard::Delete))
					{
						static std::vector<int> selected_links;
						selected_links.resize(static_cast<size_t>(num_selected));
						ImNodes::GetSelectedLinks(selected_links.data());
						for (const int edge_id : selected_links)
						{
							graph_.erase_edge(edge_id);
						}
					}
				}

				{
					const int num_selected = ImNodes::NumSelectedNodes();
					if (num_selected > 0 && ImGui::IsKeyReleased(DirectX::Keyboard::Delete))
					{
						static std::vector<int> selected_nodes;
						selected_nodes.resize(static_cast<size_t>(num_selected));
						ImNodes::GetSelectedNodes(selected_nodes.data());
						for (const int node_id : selected_nodes)
						{
							graph_.erase_node(node_id);
							auto iter = std::find_if(
								nodes_.begin(), nodes_.end(), [node_id](const UiNode& node) -> bool {
									return node.id == node_id;
								});
							// Erase any additional internal nodes
							switch (iter->type)
							{
							case UiNodeType::add:
								graph_.erase_node(iter->add.lhs);
								graph_.erase_node(iter->add.rhs);
								break;
							case UiNodeType::multiply:
								graph_.erase_node(iter->multiply.lhs);
								graph_.erase_node(iter->multiply.rhs);
								break;
							case UiNodeType::output:
								graph_.erase_node(iter->output.r);
								graph_.erase_node(iter->output.g);
								graph_.erase_node(iter->output.b);
								root_node_id_ = -1;
								break;
							case UiNodeType::sine:
								graph_.erase_node(iter->sine.input);
								break;
							default:
								break;
							}
							nodes_.erase(iter);
						}
					}
				}

				ImGui::End();

				// The color output window

				const ImU32 color =
					root_node_id_ != -1 ? evaluate(graph_, root_node_id_) : IM_COL32(255, 20, 0, 255);
				ImGui::PushStyleColor(ImGuiCol_WindowBg, color);
				ImGui::Begin("output color");
				ImGui::End();
				ImGui::PopStyleColor();
			}
		private:
			ImU32 evaluate(const Graph<Node>& graph, const int root_node);
		private:
			enum class UiNodeType
			{
				add,
				multiply,
				output,
				sine,
				time,
			};

			struct UiNode
			{
				UiNodeType type;
				// The identifying id of the ui node. For add, multiply, sine, and time
				// this is the "operation" node id. The additional input nodes are
				// stored in the structs.
				int id;

				union
				{
					struct
					{
						int lhs, rhs;
					} add;

					struct
					{
						int lhs, rhs;
					} multiply;

					struct
					{
						int r, g, b;
					} output;

					struct
					{
						int input;
					} sine;
				};
			};

			Graph<Node>            graph_;
			std::vector<UiNode>    nodes_;
			int                    root_node_id_;
			ImNodesMiniMapLocation minimap_location_;
			float current_time_seconds;
		};

		static ColorNodeEditor color_editor;

} // namespace example
